#include <catch2/catch_test_macros.hpp>
#include "MultiGrid.h"
#include "BoundaryConditions.h"

Real bottom(const Point<2> &p)
{
    return p[1];
}

Real up(const Point<2> &p)
{
    return 1 - p[1];
}

template <int N>
Real left(const Point<N> &p)
{
    return p[0];
}

template <int N>
Real right(const Point<N> &p)
{
    return 1 - p[0];
}

template <int N>
Real zero(const Point<N> &p)
{
    return 0;
}

SingleBoundaryCondition<2> Bottom(bottom, zero<2>, 0, 0);
SingleBoundaryCondition<2> Up(up, zero<2>, 0, 0);
SingleBoundaryCondition<2> Left(left<2>, zero<2>, 0, 0);
SingleBoundaryCondition<2> Right(right<2>, zero<2>, 0, 0);

TEST_CASE("test the discretization", "[GRID]")
{
    MultiGridSolver<2> MGS{GridOp::Linear_Interp<Real, 2>(), GridOp::Injection<Real, 2>()};
    MGS.define_boundary(BoundaryConditions<2>({Bottom, Up, Left, Right}));
    MGS.assignF(zero<2>);
    MGS.discretization(8);
    // MGS.PrintAdj(std::cout);
}

TEST_CASE("Weighted Jacobi Relaxation", "[RELAX]")
{
    Real mat[] = {1, 1, 1,
                  0, 1, 1,
                  0, 0, 1};
    std::vector<Real> x = {0, 0, 0};
    std::vector<Real> b = {3, 2, 1};
    std::vector<Real> e2 = {0, 0, 0};
    Real epsil = 1e-6;
    int k = 100;
    for (int i = 0; i < k; ++i)
    {
        Weighted_Jacobi_Relax(mat, 3, x.begin(), b.begin(), 2.0 / 3, 1);
        Real norm2 = 0;
        for (int j = 0; j < 3; ++j)
        {
            e2[j] = (x[j] - 1) * (x[j] - 1);
            norm2 += e2[j];
        }
        // printf("k : %d , 2-norm : %f \n", i, norm2);
        if (norm2 < epsil)
            break;
    }
}

template <int c>
Real constant(Point<1> p)
{
    return c;
}

Real my_Sin(Point<1> p)
{
    return std::sin(p[0]);
}

SingleBoundaryCondition<1> Left1d(left<1>, zero<1>, 1, 0);
SingleBoundaryCondition<1> Right1d(right<1>, constant<1>, 1, 0);

TEST_CASE("1d problem for u = x^2 ", "[SOLVER]")
{
    double epsil = 1e-6;
    int n = 8;
    GridOp::Linear_Interp<Real, 1> Linear_Interp;
    GridOp::FullWeight<Real, 1> Full_Weight;
    MultiGridSolver<1> MGS(Linear_Interp, Full_Weight);
    MGS.assignF(constant<-2>);
    MGS.define_boundary(BoundaryConditions<1>({Left1d, Right1d}));
    MGS.discretization(n);
    // MGS.PrintAdj(std::cout);
    std::vector<Real> solution(n + 1);
    // MGS.PrintMatrix(std::cout);
    double norm2 = 0;

    SECTION("use VCycle ")
    {
        std::cout << "Solving one dimension problem with VCycle method!" << std::endl;
        MGS.solve(solution,"VCycle",4,4, 1e-6, 10);
    }

    SECTION("use FMG")
    {
        std::cout << "Solving one dimension problem with FMG method!" << std::endl;
        int k = 0;
        MGS.solve(solution, "FMG",4,4);
        std::vector<Real> e(n + 1);
        for (int i = 0; i < n + 1; ++i)
        {
            e[i] = solution[i] - (Real)i * i / (n * n);
        }
        std::cout << "The 2-norm of error : " << my_vec_norm(e.data(), n + 1, 2) << std::endl;
    }

    SECTION("use direct bottom solver")
    {
        std::cout << "Changing the bottom solver to direct solver..." << std::endl;
        MGS.isDirectBottomSolver = true;
        std::cout << "Solving one dimension problem with VCycle method!" << std::endl;
        MGS.solve(solution, "VCycle",4,4, 1e-6, 10);
    }
}

